#PBAR_FD_Binning.py

import numpy as np

# PLOT SPECTRA FOR A LIST OF DETECTORS FROM A SINGLE DATASET
# groupName =  'NoSampleAir'
# groupName =  'NoSampleFe'
#groupName =  'NoSampleH2O'
groupName =  'WaterDU'
#groupName =  'IronDU'

#groupName = 'NoSampleBackground'
#detectorNumber = 35

countThreshold = 5

plt.figure()

index = datasetGroupsIndices[groupName]
#detectorList = plastic
#detectorList = liquid[0:8]
detectorList = np.arange(40)
filterSpectra = False
plotLegend = True
filterWidth = 1
minBin = 15

newBinEdges = 15 + np.arange(33) *6
ReBinSpectra = True


# Plot only those that have counts that exceed this threshold.
binThrehsoldPlotting = 185  

binList = list()
binList = np.zeros((len(datasetGroupsIndices[groupName]), len(detectorList)))

for jj in np.arange(len(datasetGroupsIndices[groupName])):
    index = datasetGroupsIndices[groupName][jj]

    for ii in np.arange(len(detectorList)):
        i = detectorList[ii]
        
        if filterSpectra:
            y = ndimage.filters.gaussian_filter(dat[index][i,:], filterWidth, order = 0)
        else:
            y = dat[index][i,:]
        # find the threshold
        temp = where((np.arange(len(y)) > minBin) & (y < countThreshold))[0]
        if len(temp) == 0:
            temp = np.array((256, 256))
        #binList.append(temp[0])
        binList[jj,ii] = temp[0]
        if temp[0] < binThrehsoldPlotting:
            continue
        plt.plot(y, label = prefixList[index] + ', ' + str(i+1) + ', bin = ' + str(temp[0]))
        #ylimit = np.array(plt.ylim())
        plt.plot(temp[0] * np.ones(2), np.array((1, 1e6)))
        
plt.title(groupName)
if plotLegend:
    plt.legend(loc = 1)
plt.show()
plt.yscale('log')


# PLOT SPECTRA FOR A LIST OF DETECTORS FROM A SINGLE DATASET
# groupName =  'NoSampleAir'
# groupName =  'NoSampleFe'
#groupName =  'NoSampleH2O'
# groupName =  'WaterDU'
#groupName =  'IronDU'
groupName = 'ActiveBackground'

#groupName = 'NoSampleBackground'
#detectorNumber = 35

countThreshold = 5

plt.figure()

index = datasetGroupsIndices[groupName]
#detectorList = plastic
#detectorList = liquid[0:8]
detectorList = np.arange(40)
filterSpectra = False

plotLegend = False
ReBinSpectra = True
filterWidth = 1
minBin = 15

newBinSpacing = 6
newBinStart = 15
newBinEdges = newBinStart + np.arange(33) * newBinSpacing

# plot only those that have counts that exceed this threshold.
binThrehsoldPlotting = 1

binList = list()
binList = np.zeros((len(datasetGroupsIndices[groupName]), len(detectorList)))

for jj in np.arange(len(datasetGroupsIndices[groupName])):
    index = datasetGroupsIndices[groupName][jj]
    
    # rebin the data
    if ReBinSpectra:
        (datTemp, newBinCenters) = PBAR_FD.ReBinData(dat[index], np.arange(256), newBinEdges)
        datTemp = datTemp/newBinSpacing
        binCenters = newBinCenters.copy()
    else:
        datTemp = dat[index]
        binCenters = np.arange(dat[index].shape[1])

    for ii in np.arange(len(detectorList)):
        i = detectorList[ii]
        
        if filterSpectra:
            y = ndimage.filters.gaussian_filter(datTemp[i,:], filterWidth, order = 0)
        else:
            y = datTemp[i,:]
            
        # find the threshold
        temp = where((np.arange(len(y)) > minBin) & (y < countThreshold))[0]
        if len(temp) == 0:
            temp = np.array((len(binCenters), len(binCenters))) -1
        binList[jj,ii] = binCenters[temp[0]]
        if binList[jj,ii] < binThrehsoldPlotting:
            continue
        plt.plot(binCenters, y, label = prefixList[index] + ', ' + str(i+1) + ', bin = ' + str(binList[jj,ii]))
        #ylimit = np.array(plt.ylim())
        
        plt.plot(binList[jj,ii] * np.ones(2), np.array((1, 1e6)))
        
plt.title(groupName)
if plotLegend:
    plt.legend(loc = 1)
plt.show()
plt.yscale('log')
plt.xlim((0, 255))


# PLOT OF CUMMULATIVE SUM VERSUS START BIN

datTemp = dat[datasetGroupsIndices['IronDU'][0]]
datBkgTemp = dat[datasetGroupsIndices['ActiveBackground'][1]]

threshold = 38

datNet = datTemp - datBkgTemp
datNet[:,0:threshold] = 0

datNetCumSum = datNet.cumsum(axis = 1)

for ii in np.arange(datNetCumSum.shape[0]):
    datNetCumSum[ii,:] = datNetCumSum[ii,:] / datNetCumSum[ii,-2]

plt.figure()
plt.plot(datNetCumSum[detGroups['plastic']].T)
plt.grid()


#  Calculate gain shift on gross binning
datSignalTemp = dat[datasetGroupsIndices['IronDU'][0]] # 'dz21'
datBkgTemp = dat[datasetGroupsIndices['ActiveBackground'][3]]  # 'dz22 - not-so-activated dataset
datBkgTemp = dat[datasetGroupsIndices['ActiveBackground'][4]]  # 'dz

newBinEdgesList  = list( (np.arange(257).astype(float), 15 + np.arange(33) *6., 15 + np.arange(33) *5.  , 15 + np.arange(33) *4.))
datNewBin = list()
datNewBin2 = list()

newBinCentersList = list()

binSizeList = np.array((1, 6, 5, 4))
# Save the original stuff
datNewBin.append(datSignalTemp - datBkgTemp)
datNewBin2.append(datSignalTemp)

newBinCentersList.append(np.arange(256) + 0.5)

# Rebin the data
for ii in np.arange(len(newBinEdgesList)-1):
    (temp, newBinCenters) = PBAR_FD.ReBinData(datSignalTemp - datBkgTemp, np.arange(256)+0.5, newBinEdgesList[ii+1])
    datNewBin.append(temp)
    newBinCentersList.append(newBinCenters)
    (temp, newBinCenters) = PBAR_FD.ReBinData(datSignalTemp, np.arange(256)+0.5, newBinEdgesList[ii+1])
    datNewBin2.append(temp)

threshold =  40
binBounds = np.array((40, 80))



# Fit exponential to the background-subtracted spectra
detectorNum = 1

pfit2 = np.zeros((len(newBinCentersList), 2))

for ii in np.arange(len(newBinCentersList)):
    cut = (newBinCentersList[ii] > binBounds[0])  & (newBinCentersList[ii] < binBounds[1])  
    x = newBinCentersList[ii][cut]
    y = log(datNewBin[ii][detectorNum, cut]/ (x[1] - x[0]))
    cutt = y > 0
    
    pfit2[ii,:] = np.polyfit(x[cutt], y[cutt], 1)
    
    


# Shift the threshold X percent and calculate the total counts for each
gainShiftList = np.linspace(0.9, 1.10, 20)
sideCountList = list()
totalCountList = list()

for gg in np.arange(len(gainShiftList)):
    gainShift = gainShiftList[gg]
    nSize = 10;
    countBinBounds = np.array((40, 256))
    fitBinBounds = np.array((37, 60))  # section to fit exponential
    
    fitBinBoundsShifted = gainShift * fitBinBounds
    countBinBoundsShifted = gainShift * countBinBounds
    
    pfit = np.zeros((len(newBinCentersList), 2))
    
    sideCount = np.zeros(len(newBinCentersList))
    totalCount = np.zeros(len(newBinCentersList))
    for ii in np.arange(len(newBinCentersList)):
        
        # fit exponential to subset
        cut = (newBinCentersList[ii] >= fitBinBoundsShifted[0])  & (newBinCentersList[ii] <= fitBinBoundsShifted[1])  
        x = newBinCentersList[ii][cut]
        y = log(datNewBin[ii][detectorNum, cut])
        cutt = y > 0
        pfit[ii,:] = np.polyfit(x[cutt], y[cutt], deg =1, w = datNewBin[ii][detectorNum, cut])
    
        # Find which binedges are contained - only  
        cut = (newBinEdgesList[ii] >= countBinBoundsShifted[0]) & (newBinEdgesList[ii] <= countBinBoundsShifted[1])
        binIndices = find(cut)
        newBinWidth = (newBinEdgesList[ii][2] - newBinEdgesList[ii][1])
        
        # Calculate the counts in the fractional bins and    
        # Calculate left fractional count
        if newBinEdgesList[ii][binIndices[0]] == countBinBoundsShifted[0]:
            sideCount[ii] = 0
        else:
            temp = np.linspace(countBinBoundsShifted[0], newBinEdgesList[ii][binIndices[0]], nSize)
            
            area1 = 0.5 * (np.polyval(pfit[ii,:], newBinEdgesList[ii][binIndices[0]-1]) + np.polyval(pfit[ii,:], newBinEdgesList[ii][binIndices[0]])) * newBinWidth
            area2 = 0.5 * (np.polyval(pfit[ii,:], countBinBoundsShifted[0]) + np.polyval(pfit[ii,:], newBinEdgesList[ii][binIndices[0]])) * \
                (newBinEdgesList[ii][binIndices[0]] - countBinBoundsShifted[0])        
            sideCount[ii] = datNewBin[ii][detectorNum,binIndices[0]-1] * area2/area1
            
            #binCenters = (temp[1:] + temp[:-1])/2
            #sideCount[ii] = sum(np.exp(np.polyval(pfit[ii,:], binCenters))) / len(binCenters) * 
        totalCount[ii] = sideCount[ii] + sum(datNewBin[ii][detectorNum,binIndices[0]:])
    sideCountList.append(sideCount)
    totalCountList.append(totalCount)

sideCountList = np.array(sideCountList)
totalCountList = np.array(totalCountList)


# Plot the specta with different binning and background subtraction

plt.figure()
ii = 0
plot(newBinCentersList[0], datSignalTemp[detectorNum, :] / (newBinCentersList[ii][1] - newBinCentersList[ii][0]), label = 'DU')
plot(newBinCentersList[0], datBkgTemp[detectorNum, :] / (newBinCentersList[ii][1] - newBinCentersList[ii][0]), label = 'Active Background')


#for ii in np.arange(len(datNewBin)):
    #plot(newBinCentersList[ii], \
         #datNewBin[ii][detectorNum, :] / (newBinCentersList[ii][1] - newBinCentersList[ii][0]), '.-', label =  'DU - AB, bin size = ' +str(binSizeList[ii]))

for ii in np.arange(len(datNewBin)):
    plot(newBinCentersList[ii], \
         datNewBin2[ii][detectorNum, :] / (newBinCentersList[ii][1] - newBinCentersList[ii][0]), '.-', label =  'DU, bin size = ' +str(binSizeList[ii]), markersize = 10)


# plot fits
xarray = np.arange(255)
#for ii in np.arange(len(datNewBin)):
for ii in np.arange(1):
    plot(xarray, np.exp(np.polyval(pfit[ii,:], xarray)))

legend()
plt.xlabel('Bin')
plt.ylabel('Counts/Bin')
plt.grid()
plt.xlim((0, 150))
plt.yscale('log')



# Plot the specta with different binning and background subtraction

plt.figure()
ii = 0
plot(newBinCentersList[0], datSignalTemp[detectorNum, :] / (newBinCentersList[ii][1] - newBinCentersList[ii][0]), label = 'DU')
plot(newBinCentersList[0], datBkgTemp[detectorNum, :] / (newBinCentersList[ii][1] - newBinCentersList[ii][0]), label = 'Active Background')


for ii in np.arange(len(datNewBin)):
    plot(newBinCentersList[ii], \
         datNewBin[ii][detectorNum, :] / (newBinCentersList[ii][1] - newBinCentersList[ii][0]), \
         '.-', label =  'DU - AB, bin size = ' +str(binSizeList[ii]))

# plot fits
xarray = np.arange(255)
#for ii in np.arange(len(datNewBin)):
for ii in np.arange(1):
    plot(xarray, np.exp(np.polyval(pfit[ii,:], xarray)))

legend()
plt.xlabel('Bin')
plt.ylabel('Counts/Bin')
plt.grid()
plt.xlim((0, 150))
plt.yscale('log')


#  Plot counts from a threshold
plt.figure()
plt.grid()
for ii in np.arange(len(binSizeList)):
    plot(gainShiftList, totalCountList[:,ii], label = 'bin size ' + str(binSizeList[ii]))
plt.xlabel('Shift in threshold')
plt.ylabel('Counts')
plt.legend()

# plot of the fits

figure()

plot(newBinCentersList[ii], datNewBin[ii][detectorNum,:])
plot(newBinCentersList[ii][binIndices[0]:], datNewBin[ii][detectorNum,binIndices[0]:])    
    